from typing import TYPE_CHECKING, Union

import torch
from zkl_metrics import MeanMetric, Metric
from zkl_training import ProcessingTaskPlugin

if TYPE_CHECKING:
    from .training import GPTTraining, GPTTrainingOutput
    from .validating import GPTValidating


class MetricsPlugin(ProcessingTaskPlugin):
    def __init__(self):
        super().__init__()
        self.ce_metric = MeanMetric()
        self.acc_metric = MeanMetric()

    @property
    def task(self) -> Union['GPTTraining', 'GPTValidating']:
        task = super().task
        from .training import GPTTraining
        from .validating import GPTValidating
        assert isinstance(task, (GPTTraining, GPTValidating))
        return task

    # task

    def on_after_process(self, output: 'GPTTrainingOutput'):
        batch_tokens_ce = torch.asarray(output.tokens_out_ce.detach(), dtype=torch.float64)
        batch_tokens_acc = torch.asarray(output.tokens_out_acc.detach(), dtype=torch.float64)
        batch_tokens_mask = torch.asarray(output.tokens_out_mask, dtype=torch.bool)
        # [..., chunk_size]

        ce = torch.masked.sum(batch_tokens_ce, mask=batch_tokens_mask)
        acc = torch.masked.sum(batch_tokens_acc, mask=batch_tokens_mask)
        weight = torch.sum(batch_tokens_mask, dtype=torch.float64)
        # []

        self.ce_metric.update(ce, weight)
        self.acc_metric.update(acc, weight)

    # metrics

    @property
    def metrics(self) -> dict[str, Metric]:
        return {
            attr_name[:-len('_metric')]: attr_value
            for attr_name, attr_value in vars(self).items()
            if attr_name.endswith('_metric') and isinstance(attr_value, Metric)}
